// DXGI + D3D11 Desktop Duplication -> NV12 GPU conversion -> save raw NV12
// Compile with: d3d11.lib dxgi.lib

#include <windows.h>
#include <d3d11.h>
#include <dxgi1_2.h>
#include <wrl/client.h>
#include <vector>
#include <fstream>
#include <string>
#include <stdio.h>
#include <iostream>

// Media Foundation
#include <mfapi.h>
#include <mfplay.h>
#include <mfreadwrite.h>
#include <mferror.h>
#include <Codecapi.h>

// Simplified error handling via macros
#ifdef DEBUG
#define ON_ERROR __debugbreak(); return 1;
#else
#define ON_ERROR system("pause"); return 1;
#endif

#define CHECK(x, err) if (!(x)) { std::cerr << err << std::endl; ON_ERROR; }
#define CHECK_HR(x, err) if (FAILED(x)) { std::cerr << err << std::endl; ON_ERROR; }

using namespace Microsoft::WRL;

#pragma comment(lib, "d3d11.lib")
#pragma comment(lib, "Mfplat.lib")
#pragma comment(lib, "mfuuid.lib")

ComPtr<ID3D11Device> gDevice;
ComPtr<ID3D11DeviceContext> gContext;
ComPtr<ID3D11VideoDevice> gVideoDevice;
ComPtr<ID3D11VideoContext> gVideoContext;
ComPtr<ID3D11VideoProcessorEnumerator> gVPEnum;
ComPtr<ID3D11VideoProcessor> gVP;
ComPtr<ID3D11VideoProcessorOutputView> gOutputView;
ComPtr<ID3D11Texture2D> gNV12Tex;

bool InitD3D11() {
    D3D_FEATURE_LEVEL level;
    if (FAILED(D3D11CreateDevice(
        nullptr, D3D_DRIVER_TYPE_HARDWARE, nullptr,
        D3D11_CREATE_DEVICE_BGRA_SUPPORT, nullptr, 0,
        D3D11_SDK_VERSION, &gDevice, &level, &gContext))) return false;

    gDevice.As(&gVideoDevice);
    gContext.As(&gVideoContext);
    return true;
}

ComPtr<IDXGIOutputDuplication> InitDuplication(UINT& width, UINT& height) {
    ComPtr<IDXGIDevice> dxgiDevice;
    gDevice.As(&dxgiDevice);

    ComPtr<IDXGIAdapter> adapter;
    dxgiDevice->GetAdapter(&adapter);

    ComPtr<IDXGIOutput> output;
    adapter->EnumOutputs(0, &output);

    DXGI_OUTPUT_DESC desc;
    output->GetDesc(&desc);

    ComPtr<IDXGIOutput1> output1;
    output.As(&output1);

    DXGI_OUTDUPL_DESC duplDesc;
    ComPtr<IDXGIOutputDuplication> duplication;
    output1->DuplicateOutput(gDevice.Get(), &duplication);
    duplication->GetDesc(&duplDesc);

    width = duplDesc.ModeDesc.Width;
    height = duplDesc.ModeDesc.Height;

    return duplication;
}

bool InitVideoProcessor(UINT width, UINT height) {
    D3D11_VIDEO_PROCESSOR_CONTENT_DESC desc = {};
    desc.InputFrameFormat = D3D11_VIDEO_FRAME_FORMAT_PROGRESSIVE;
    desc.InputWidth = width;
    desc.InputHeight = height;
    desc.OutputWidth = width;
    desc.OutputHeight = height;
    desc.Usage = D3D11_VIDEO_USAGE_PLAYBACK_NORMAL;

    if (FAILED(gVideoDevice->CreateVideoProcessorEnumerator(&desc, &gVPEnum))) return false;
    if (FAILED(gVideoDevice->CreateVideoProcessor(gVPEnum.Get(), 0, &gVP))) return false;

    D3D11_TEXTURE2D_DESC texDesc = {};
    texDesc.Width = width;
    texDesc.Height = height;
    texDesc.MipLevels = 1;
    texDesc.ArraySize = 1;
    texDesc.Format = DXGI_FORMAT_NV12;
    texDesc.SampleDesc.Count = 1;
    texDesc.Usage = D3D11_USAGE_DEFAULT;
    texDesc.BindFlags = D3D11_BIND_RENDER_TARGET | D3D11_BIND_SHADER_RESOURCE;

    if (FAILED(gDevice->CreateTexture2D(&texDesc, nullptr, &gNV12Tex))) return false;

    D3D11_VIDEO_PROCESSOR_OUTPUT_VIEW_DESC ovDesc = {};
    ovDesc.ViewDimension = D3D11_VPOV_DIMENSION_TEXTURE2D;
    ovDesc.Texture2D.MipSlice = 0;
    if (FAILED(gVideoDevice->CreateVideoProcessorOutputView(
        gNV12Tex.Get(), gVPEnum.Get(), &ovDesc, &gOutputView))) return false;

    return true;
}

int main() {
    HRESULT hr;

    if (!InitD3D11()) return -1;

    UINT width = 0, height = 0;
    auto duplication = InitDuplication(width, height);
    if (!duplication) return -1;

    if (!InitVideoProcessor(width, height)) return -1;

    ComPtr<IMFDXGIDeviceManager> deviceManager;
    ComPtr<IMFTransform> transform;
    ComPtr<IMFAttributes> attributes;
    ComPtr<IMFMediaEventGenerator> eventGen;
    DWORD inputStreamID;
    DWORD outputStreamID;

    // ------------------------------------------------------------------------
    // Initialize Media Foundation & COM
    // ------------------------------------------------------------------------

    hr = MFStartup(MF_VERSION);
    CHECK_HR(hr, "Failed to start Media Foundation");

    // Create device manager
    UINT resetToken;
    hr = MFCreateDXGIDeviceManager(&resetToken, &deviceManager);
    CHECK_HR(hr, "Failed to create DXGIDeviceManager");

    hr = deviceManager->ResetDevice(gDevice.Get(), resetToken);
    CHECK_HR(hr, "Failed to assign D3D device to device manager");

    // ------------------------------------------------------------------------
    // Initialize hardware encoder MFT
    // ------------------------------------------------------------------------

    // Find encoder
    IMFActivate ** activateRaw;
    UINT32 activateCount = 0;

    // h264 output
    MFT_REGISTER_TYPE_INFO info = { MFMediaType_Video, MFVideoFormat_H264 };

    hr = MFTEnumEx(
        MFT_CATEGORY_VIDEO_ENCODER,
        MFT_ENUM_FLAG_HARDWARE | MFT_ENUM_FLAG_SORTANDFILTER,
        NULL,
        &info,
        &activateRaw,
        &activateCount
    );
    CHECK_HR(hr, "Failed to enumerate MFTs");

    CHECK(activateCount, "No MFTs found");

    // Choose the first available encoder
    ComPtr<IMFActivate> activate = activateRaw[0];

    for (UINT32 i = 0; i < activateCount; i++)
        activateRaw[i]->Release();

    // Activate
    hr = activate->ActivateObject(IID_PPV_ARGS(&transform));
    CHECK_HR(hr, "Failed to activate MFT");

    // Get attributes
    hr = transform->GetAttributes(&attributes);
    CHECK_HR(hr, "Failed to get MFT attributes");

    // Get encoder name
    UINT32 nameLength = 0;
    std::wstring name;

    hr = attributes->GetStringLength(MFT_FRIENDLY_NAME_Attribute, &nameLength);
    CHECK_HR(hr, "Failed to get MFT name length");

    // IMFAttributes::GetString returns a null-terminated wide string
    name.resize(nameLength + 1);

    hr = attributes->GetString(MFT_FRIENDLY_NAME_Attribute, &name[0], name.size(), &nameLength);
    CHECK_HR(hr, "Failed to get MFT name");

    name.resize(nameLength);

    std::wcout << name << std::endl;

    // Unlock the transform for async use and get event generator
    hr = attributes->SetUINT32(MF_TRANSFORM_ASYNC_UNLOCK, TRUE);
    CHECK_HR(hr, "Failed to unlock MFT");

    transform.As(&eventGen);
    CHECK(eventGen, "Failed to QI for event generator");

    // Get stream IDs (expect 1 input and 1 output stream)
    hr = transform->GetStreamIDs(1, &inputStreamID, 1, &outputStreamID);
    if (hr == E_NOTIMPL)
    {
        inputStreamID = 0;
        outputStreamID = 0;
        hr = S_OK;
    }
    CHECK_HR(hr, "Failed to get stream IDs");

    // ------------------------------------------------------------------------
    // Configure hardware encoder MFT
    // ------------------------------------------------------------------------

    // Set D3D manager
    hr = transform->ProcessMessage(MFT_MESSAGE_SET_D3D_MANAGER, reinterpret_cast<ULONG_PTR>(deviceManager.Get()));
    CHECK_HR(hr, "Failed to set D3D manager");

    // Set low latency hint
    hr = attributes->SetUINT32(MF_LOW_LATENCY, TRUE);
    CHECK_HR(hr, "Failed to set MF_LOW_LATENCY");

    // Set output type
    ComPtr<IMFMediaType> outputType;

    hr = MFCreateMediaType(&outputType);
    CHECK_HR(hr, "Failed to create media type");

    hr = outputType->SetGUID(MF_MT_MAJOR_TYPE, MFMediaType_Video);
    CHECK_HR(hr, "Failed to set MF_MT_MAJOR_TYPE on H264 output media type");

    hr = outputType->SetGUID(MF_MT_SUBTYPE, MFVideoFormat_H264);
    CHECK_HR(hr, "Failed to set MF_MT_SUBTYPE on H264 output media type");

    hr = outputType->SetUINT32(MF_MT_AVG_BITRATE, 30000000);
    CHECK_HR(hr, "Failed to set average bit rate on H264 output media type");

    hr = MFSetAttributeSize(outputType.Get(), MF_MT_FRAME_SIZE, width, height);
    CHECK_HR(hr, "Failed to set frame size on H264 MFT out type");

    hr = MFSetAttributeRatio(outputType.Get(), MF_MT_FRAME_RATE, 60, 1);
    CHECK_HR(hr, "Failed to set frame rate on H264 MFT out type");

    hr = outputType->SetUINT32(MF_MT_INTERLACE_MODE, 2);
    CHECK_HR(hr, "Failed to set MF_MT_INTERLACE_MODE on H.264 encoder MFT");

    hr = outputType->SetUINT32(MF_MT_ALL_SAMPLES_INDEPENDENT, TRUE);
    CHECK_HR(hr, "Failed to set MF_MT_ALL_SAMPLES_INDEPENDENT on H.264 encoder MFT");

    hr = transform->SetOutputType(outputStreamID, outputType.Get(), 0);
    CHECK_HR(hr, "Failed to set output media type on H.264 encoder MFT");

    MFT_OUTPUT_STREAM_INFO o_info;
    hr = transform->GetOutputStreamInfo(outputStreamID, &o_info);
    CHECK_HR(hr, "Failed to get output info on H.264 encoder MFT");

    // Set input type
    ComPtr<IMFMediaType> inputType;

    hr = MFCreateMediaType(&inputType);
    CHECK_HR(hr, "Failed to create media type");

    for (DWORD i = 0;; i++)
    {
        inputType = nullptr;
        hr = transform->GetInputAvailableType(inputStreamID, i, &inputType);
        CHECK_HR(hr, "Failed to get input type");

        hr = inputType->SetGUID(MF_MT_MAJOR_TYPE, MFMediaType_Video);
        CHECK_HR(hr, "Failed to set MF_MT_MAJOR_TYPE on H264 MFT input type");

        hr = inputType->SetGUID(MF_MT_SUBTYPE, MFVideoFormat_NV12);
        CHECK_HR(hr, "Failed to set MF_MT_SUBTYPE on H264 MFT input type");

        hr = MFSetAttributeSize(inputType.Get(), MF_MT_FRAME_SIZE, width, height);
        CHECK_HR(hr, "Failed to set MF_MT_FRAME_SIZE on H264 MFT input type");

        hr = MFSetAttributeRatio(inputType.Get(), MF_MT_FRAME_RATE, 60, 1);
        CHECK_HR(hr, "Failed to set MF_MT_FRAME_RATE on H264 MFT input type");

        hr = transform->SetInputType(inputStreamID, inputType.Get(), 0);
        CHECK_HR(hr, "Failed to set input type");

        break;
    }

    // ------------------------------------------------------------------------
    // Start encoding
    // ------------------------------------------------------------------------
    hr = transform->ProcessMessage(MFT_MESSAGE_COMMAND_FLUSH, NULL);
    CHECK_HR(hr, "Failed to process FLUSH command on H.264 MFT");

    hr = transform->ProcessMessage(MFT_MESSAGE_NOTIFY_BEGIN_STREAMING, NULL);
    CHECK_HR(hr, "Failed to process BEGIN_STREAMING command on H.264 MFT");

    hr = transform->ProcessMessage(MFT_MESSAGE_NOTIFY_START_OF_STREAM, NULL);
    CHECK_HR(hr, "Failed to process START_OF_STREAM command on H.264 MFT");

    std::ofstream fout("output.h264", std::ios::binary);

    DXGI_OUTDUPL_FRAME_INFO frameInfo = {};
    ComPtr<IDXGIResource> desktopResource;
    ComPtr<ID3D11Texture2D> acquiredTex;

    int frame_num = 0;
    while(frame_num < 600) {
        ComPtr<IMFMediaEvent> event;
        hr = eventGen->GetEvent(0, &event);
        CHECK_HR(hr, "Failed to get next event");

        MediaEventType eventType;
        hr = event->GetType(&eventType);
        CHECK_HR(hr, "Failed to get event type");

        switch (eventType)
        {
        case METransformNeedInput: {
            HRESULT hr = duplication->AcquireNextFrame(100, &frameInfo, &desktopResource);
            if (FAILED(hr)) continue;

            desktopResource.As(&acquiredTex);

            ComPtr<ID3D11VideoProcessorInputView> inputView;
            D3D11_VIDEO_PROCESSOR_INPUT_VIEW_DESC ivDesc = {};
            ivDesc.ViewDimension = D3D11_VPIV_DIMENSION_TEXTURE2D;
            ivDesc.Texture2D.ArraySlice = 0;

            gVideoDevice->CreateVideoProcessorInputView(acquiredTex.Get(), gVPEnum.Get(), &ivDesc, &inputView);

            D3D11_VIDEO_PROCESSOR_STREAM stream = {};
            stream.Enable = TRUE;
            stream.pInputSurface = inputView.Get();

            gVideoContext->VideoProcessorBlt(gVP.Get(), gOutputView.Get(), 0, 1, &stream);

            // Create buffer
            ComPtr<IMFMediaBuffer> inputBuffer;
            hr = MFCreateDXGISurfaceBuffer(__uuidof(ID3D11Texture2D), gNV12Tex.Get(), 0, FALSE, &inputBuffer);
            CHECK_HR(hr, "Failed to create IMFMediaBuffer");

            // Create sample
            ComPtr<IMFSample> sample;
            hr = MFCreateSample(&sample);
            CHECK_HR(hr, "Failed to create IMFSample");
            hr = sample->AddBuffer(inputBuffer.Get());
            CHECK_HR(hr, "Failed to add buffer to IMFSample");

            hr = transform->ProcessInput(inputStreamID, sample.Get(), 0);
            CHECK_HR(hr, "Failed to process input");

            Sleep(33);
        }
        break;

        case METransformHaveOutput: {
            DWORD status;
            MFT_OUTPUT_DATA_BUFFER outputBuffer;
            outputBuffer.dwStreamID = outputStreamID;
            outputBuffer.pSample = nullptr;
            outputBuffer.dwStatus = 0;
            outputBuffer.pEvents = nullptr;

            hr = transform->ProcessOutput(0, 1, &outputBuffer, &status);
            CHECK_HR(hr, "ProcessOutput failed");

            // Release sample as it is not processed any further.
            if (outputBuffer.pSample) {
                ComPtr<IMFMediaBuffer> oBuffer;
                hr = outputBuffer.pSample->ConvertToContiguousBuffer(&oBuffer);
                if (SUCCEEDED(hr)) {
                    BYTE* outData = nullptr;
                    DWORD maxLen = 0, currLen = 0;
                    oBuffer->Lock(&outData, &maxLen, &currLen);
                    fout.write((char*)outData, currLen);
                    oBuffer->Unlock();
                    frame_num++;
                    std::wcout << frame_num << std::endl;
                }
                outputBuffer.pSample->Release();
            }

            duplication->ReleaseFrame();
        }
        break;

        default:
            CHECK(false, "Unknown event");
            break;
        }
    }

    fout.close();

    // ------------------------------------------------------------------------
    // Finish encoding
    // ------------------------------------------------------------------------
    hr = transform->ProcessMessage(MFT_MESSAGE_NOTIFY_END_OF_STREAM, NULL);
    CHECK_HR(hr, "Failed to process END_OF_STREAM command on H.264 MFT");

    hr = transform->ProcessMessage(MFT_MESSAGE_NOTIFY_END_STREAMING, NULL);
    CHECK_HR(hr, "Failed to process END_STREAMING command on H.264 MFT");

    hr = transform->ProcessMessage(MFT_MESSAGE_COMMAND_FLUSH, NULL);
    CHECK_HR(hr, "Failed to process FLUSH command on H.264 MFT");
    return 0;
}